`include "define.sv"
module systolic_system (
    input clk,
    input reset,
    input [31:0] a_din [0:31],   // 32个通道的输入数据（非压缩数组）
    input a_wr_en [0:31],         // 32个通道的写使能
    input [31:0] b_din [0:31],    // 32个通道的输入数据
    input b_wr_en [0:31],         // 32个通道的写使能
    input [31:0] c_din [0:15][0:15], // 32x32矩阵输入
    input c_wr_en [0:15][0:15],      // 32x32写使能
    input [1:0] matrix_type,      // 控制信号
    input [1:0] data_type,
    input mix_precision,
    output array_ready,           // 向上级发送空闲信号
    output array_done,           // 向下级发送阵列完成信号
    output reg [31:0] c_out_muxed, // 复用后的输出
    output reg c_out_valid, // 新增：输出数据有效信号
    input wire axi_ready // 新增：AXI写通道握手信号
);

// FIFO输出连接
wire [31:0] a_in_fifo [0:31];
wire [31:0] b_in_fifo [0:31];
//wire [31:0] c_in_fifo [0:31][0:31];

reg [7:0] counter;          
reg [31:0] c_out [0:15][0:15];  // 16x16 结果存储
wire [31:0] systolic_dout [0:15][0:15];  // 来自 systolic 阵列的输出

reg adder_valid [0:15][0:15];
reg add_valid_o [0:15][0:15]; 
reg calc_done [0:15][0:15];

//---------------- 检测FIFO空状态 ------------------
wire [31:0] a_fifo_empty; // A通道FIFO空信号，32位向量
wire [31:0] b_fifo_empty; // B通道FIFO空信号，32位向量
//wire c_fifo_empty [0:31][0:31];         // C通道FIFO空信号

// FIFO空状态寄存一拍
reg [31:0] a_fifo_empty_reg; // A通道FIFO空信号寄存
reg [31:0] b_fifo_empty_reg; // B通道FIFO空信号寄存
wire [31:0] a_fifo_full; // A通道FIFO满信号，32位向量
wire [31:0] b_fifo_full; // B通道FIFO满信号，32位向量
//wire [31:0][31:0] c_fifo_full; // C通道FIFO满信号


// FIFO读使能控制
reg [31:0] a_rd_en;
reg [31:0] b_rd_en;
//reg [31:0][31:0] c_rd_en;



// 实例化32个A通道FIFO
genvar i1, j1;
generate
for (i1 = 0; i1 < 32; i1 = i1 + 1) begin: a_fifo_inst
    sync_fifo #(
       .DATA_WIDTH(32),
       .DEPTH(16)
    ) u_a_fifo (
       .clk(clk),
       .reset(reset),
       .wr_en(a_wr_en[i1]),
       .rd_en(a_rd_en[i1]),  // 使用可控的读使能
       .din(a_din[i1]),
       .dout(a_in_fifo[i1]),
       .full(a_fifo_full[i1]),
       .empty(a_fifo_empty[i1])
    );
end
endgenerate

// 实例化32个B通道FIFO
generate
for (i1 = 0; i1 < 32; i1 = i1 + 1) begin: b_fifo_inst
    sync_fifo #(
       .DATA_WIDTH(32),
       .DEPTH(16)
    ) u_b_fifo (
       .clk(clk),
       .reset(reset),
       .wr_en(b_wr_en[i1]),
       .rd_en(b_rd_en[i1]),  // 使用可控的读使能
       .din(b_din[i1]),
       .dout(b_in_fifo[i1]),
       .full(b_fifo_full[i1]),
       .empty(b_fifo_empty[i1])
    );
end
endgenerate

/* // 实例化32x32个C通道FIFO
generate
for (i1 = 0; i1 < 32; i1 = i1 + 1) begin: c_fifo_row
    for (j1 = 0; j1 < 32; j1 = j1 + 1) begin: c_fifo_col
        sync_fifo #(
           .DATA_WIDTH(32),
           .DEPTH(1)
        ) u_c_fifo (
           .clk(clk),
           .reset(reset),
           .wr_en(c_wr_en[i1][j1]),
           .rd_en(c_rd_en[i1][j1]),  // 使用可控的读使能
           .din(c_din[i1][j1]),
           .dout(c_in_fifo[i1][j1]),
           .full(c_fifo_full[i1][j1]),
           .empty(c_fifo_empty[i1][j1])
        );
    end
end
endgenerate */

// FIFO读使能控制逻辑和空状态寄存
always @(posedge clk or posedge reset) begin
    if (reset) begin
        a_rd_en <= 32'h0;
        b_rd_en <= 32'h0;
        a_fifo_empty_reg <= 32'h0;
        b_fifo_empty_reg <= 32'h0;
        /* for (int i = 0; i < 32; i++) begin
            for (int j = 0; j < 32; j++) begin
                c_rd_en[i][j] <= 1'b0;
            end
        end */
    end else begin
        // 寄存FIFO空状态
        a_fifo_empty_reg <= a_fifo_empty;
        b_fifo_empty_reg <= b_fifo_empty;
        
        // A和B FIFO的读使能：当有写操作时使能
        for (int i = 0; i < 32; i++) begin
            a_rd_en[i] <= a_wr_en[i];
            b_rd_en[i] <= b_wr_en[i];
        end
        
        /* // C FIFO的读使能：当有写操作时使能，并保持一个时钟周期
        for (int i = 0; i < 32; i++) begin
            for (int j = 0; j < 32; j++) begin
                if (c_wr_en[i][j]) begin
                    c_rd_en[i][j] <= 1'b1;
                end else begin
                    c_rd_en[i][j] <= 1'b0;
                end
            end
        end */
    end
end

// 状态机定义
typedef enum logic [1:0] {
    IDLE = 2'b00,      // 等待数据输入
    CALCULATE = 2'b01, // 数据计算中
    OUTPUT = 2'b10,    // 计算完成，正在输出数据
    DONE = 2'b11       // 输出完成，等待复位
} array_state_t;

array_state_t array_state;
array_state_t array_state_next;

// PE复位控制信号
reg pe_reset;
reg [2:0] pe_reset_counter;
reg output_complete;  // 输出完成标志

// 状态机组合逻辑
always @(*) begin
    array_state_next = array_state;
    
    case (array_state)
        IDLE: begin
            // 检测是否有C矩阵写入操作
            if (c_wr_en[0][0]) begin    //todo:需要结合fifo_to_array模块信号 FIFO存满再进入计算状态
                array_state_next = CALCULATE;
            end
        end
        
        CALCULATE: begin
            // 等待计算完成
            if (calc_done[15][15]) begin
                array_state_next = OUTPUT;
            end
        end
        
        OUTPUT: begin
            // 等待输出完成
            if (output_complete) begin
                array_state_next = DONE;
            end
        end
        
        DONE: begin
            // 等待PE复位完成
            if (pe_reset) begin
                array_state_next = IDLE;
            end
        end
        
        default: begin
            array_state_next = IDLE;
        end
    endcase
end

// 状态机时序逻辑
always @(posedge clk or posedge reset) begin
    if (reset) begin
        array_state <= IDLE;
    end else begin
        array_state <= array_state_next;
    end
end

// 输出完成检测逻辑
always @(posedge clk or posedge reset) begin
    if (reset) begin
        output_complete <= 1'b0;
    end else begin
        if (array_state == OUTPUT && !output_complete) begin
            // 开始计数输出元素
            if (counter == 8'd255) begin
                output_complete <= 1'b1;  // 输出完成
            end
        end else if (array_state == IDLE) begin
            output_complete <= 1'b0;  // 复位后清除输出完成标志
        end
    end
end

assign array_ready = (array_state == IDLE) ;    //当阵列处于等待状态时，向上级发送空闲标志信号
assign array_done = (array_state == OUTPUT && counter == 0 && axi_ready);
// 确保counter在进入OUTPUT状态时从0开始
always @(posedge clk or posedge reset) begin
    if (reset) begin
        counter <= 10'd0;
    end else if (array_state == OUTPUT) begin
        if (axi_ready) begin
            if (counter == 8'd255) begin
                counter <= 8'd0;
            end else begin
                counter <= counter + 8'd1;
            end
        end
    end else begin
        counter <= 8'd0;
    end
end

// PE复位控制逻辑
always @(posedge clk or posedge reset) begin
    if (reset) begin
        pe_reset <= 1'b0;
        pe_reset_counter <= 8'd0;
    end else begin
        // 当进入DONE状态时，延迟几个时钟周期后产生复位信号
        if (array_state == DONE && !pe_reset) begin
            pe_reset_counter <= pe_reset_counter + 3'd1;
            if (pe_reset_counter >= 3'd3) begin  // 延迟3个时钟周期
                pe_reset <= 1'b1;
                pe_reset_counter <= 3'd0;
            end
        end else if (pe_reset) begin
            pe_reset <= 1'b0;  // 复位信号持续一个时钟周期
        end
    end
end

/* // 状态机调试输出
always @(posedge clk) begin
    if (array_state != array_state_next) begin
        $display("[%t] Array State Change: %s -> %s", $time, 
                 get_state_name(array_state), get_state_name(array_state_next));
    end
end

// 状态名称函数
function string get_state_name(input [1:0] state);
    case (state)
        IDLE: return "IDLE";
        CALCULATE: return "CALCULATE";
        OUTPUT: return "OUTPUT";
        DONE: return "DONE";
        default: return "UNKNOWN";
    endcase
endfunction */

// 实例化 systolic array
systolic_array_unified u_systolic (
   .clk(clk),
   .reset(reset),
   .matrix_type(matrix_type),
   .a_in(a_in_fifo),
   .b_in(b_in_fifo),
   .c_in(c_din),
   .data_type(data_type),
   .mix_precision(mix_precision),
   .a_fifo_empty(a_fifo_empty_reg),
   .b_fifo_empty(b_fifo_empty_reg),
   .c_wr_en(c_wr_en),
   .c_out(systolic_dout),  // 连接到中间变量
   .calc_done(calc_done),
   .pe_reset(pe_reset)      // 连接PE复位信号
);

//---------------- 输出存储与复用控制 ------------------
always @(posedge clk or posedge reset) begin
    if (reset) begin
        // 复位时初始化所有元素为0
        for (int i = 0; i < 16; i++) begin
            for (int j = 0; j < 16; j++) begin
                c_out[i][j] <= 32'd0;
            end
        end
    end else begin
        // 正常更新数据
        for (int i = 0; i < 16; i++) begin
            for (int j = 0; j < 16; j++) begin
                c_out[i][j] <= systolic_dout[i][j];
            end
        end
    end
end

 //---------------- 行优先分时复用控制 ------------------

always @(posedge clk or posedge reset) begin
    if (reset) begin
        c_out_muxed <= 32'd0;
    end else if(array_state == OUTPUT) begin
        case (matrix_type)
            `m16n16k16: begin
                // 16x16矩阵：直接使用counter作为索引
                c_out_muxed <= c_out[counter[7:4]][counter[3:0]];
            end
            `m8n32k16: begin
                // 8x32矩阵：需要重新映射索引
                // 对于8x32矩阵，PE阵列的映射关系：
                // pe[0][0]->pe[7][7] 构成左半边（第1-16列）
                // pe[8][0]->pe[15][15] 构成右半边（第17-32列）
                // counter[7:4]作为行索引，counter[3:0]作为列索引
                if (counter[4:0] < 5'd16) begin
                    // 左半边：pe[0][0]->pe[7][15] -> 输出矩阵的[0:7][0:15]
                    c_out_muxed <= c_out[counter[7:5]][counter[4:0]];
                end 
                else if(counter[4:0] >= 5'd16)begin
                    // 右半边：pe[8][0]->pe[15][15] -> 输出矩阵的[0:7][16:31]
                    c_out_muxed <= c_out[counter[7:5] + 5'd8][counter[4:0] - 5'd16];
                end
            end
            `m32n8k16: begin
                // 32x8矩阵：需要重新映射索引
                // 对于32x8矩阵，PE阵列的映射关系：
                // pe[0][0]->pe[15][7] 构成上半边（第1-16行）
                // pe[16][0]->pe[31][7] 构成下半边（第17-32行）
                // counter[7:4]作为行索引，counter[3:0]作为列索引
                if (counter[7:3] < 5'd16) begin
                    // 上半边：pe[0][0]->pe[15][7] -> 输出矩阵的[0:15][0:7]
                    c_out_muxed <= c_out[counter[7:3]][counter[2:0]];
                end 
                else if(counter[7:3] >= 5'd16) begin
                    // 下半边：pe[16][0]->pe[31][7] -> 输出矩阵的[16:31][0:7]
                    c_out_muxed <= c_out[counter[7:3] - 5'd16][counter[2:0] + 4'd8];
                end
            end
            default: begin
                c_out_muxed <= 32'd0;
            end
        endcase
    end
end

always @(posedge clk or posedge reset) begin
    if (reset) begin
        c_out_valid <= 1'b0;
    end else if (array_state == OUTPUT && axi_ready) begin
        if (counter < 8'd255) begin
            c_out_valid <= 1'b1;
        end else begin
            c_out_valid <= 1'b0;
        end
    end else begin
        c_out_valid <= 1'b0;
    end
end

/* 
// 添加输出控制信号
reg output_en;
reg [7:0] elem_index;
wire [3:0]row,col;
            assign row = elem_index[7:4];  // 行索引 0-15
            assign col = elem_index[3:0];  // 列索引 0-15
//---------------- 行优先分时复用控制 ------------------
always @(posedge clk or posedge reset) begin
    if (reset) begin
        elem_index <= 8'd0;
        c_out_muxed <= 32'd0;
        output_en <= 1'b0;
    end else begin
        // 当最后一个PE计算完成时启动输出
        if (calc_done[15][15] && !output_en) begin
            output_en <= 1'b1;
            elem_index <= 8'd0;
        end
        
        // 输出使能期间进行分时复用
        if (output_en) begin
            // 计算行列索引（保持与PE阵列相同的索引）

            
            // 输出对应位置的元素
            c_out_muxed <= c_out[row][col];
            
            // 移动到下一个元素
            elem_index <= elem_index + 8'd1;
            
            // 完成256个元素输出后停止
            if (elem_index == 8'd255) begin
                output_en <= 1'b0;
            end
        end
    end 
end */

endmodule 